package leetcode.code2316;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import leetcode.IDebug;
import leetcode.helper.H;

public class Solution implements IDebug {

	class UnionFind {
		int[] roots;
		int[] sizes;
		int[] stack;

		public UnionFind(int n) {
			roots = new int[n];
			sizes = new int[n];
			for (int i = 0; i < n; i++) {
				roots[i] = i;
				sizes[i] = 1;
			}
			stack = new int[n];
		}

		public int findRoot(int id) {
			int root = 0;
			int p = 0;
			while ((root = roots[id]) != id) {
				stack[p++] = id;
				id = root;
			}
			while (p > 0) {
				roots[stack[--p]] = root;
			}
			return root;
		}

		public boolean isSameSet(int id1, int id2) {
			return findRoot(id1) == findRoot(id2);
		}

		public void union(int id1, int id2) {
			int p1 = this.findRoot(id1);
			int p2 = this.findRoot(id2);
			if (p1 != p2) {
				int size1 = this.sizes[p1];
				int size2 = this.sizes[p2];
				if (size1 < size2) {
					this.roots[p1] = p2;
					this.sizes[p2] = size1 + size2;
					this.sizes[p1] = 0;
				} else {
					this.roots[p2] = p1;
					this.sizes[p1] = size1 + size2;
					this.sizes[p2] = 0;
				}
			}
		}
	}

	public long countPairs(int n, int[][] edges) {
		UnionFind union = new UnionFind(n);
		for (int i = 0; i < edges.length; i++) {
			union.union(edges[i][0], edges[i][1]);
		}
		Map<Integer, Integer> map = new HashMap<>();
		for (int size : union.sizes) {
			if (size != 0) {
				int cnt = map.getOrDefault(size, 0) + 1;
				map.put(size, cnt);
			}
		}
		List<Integer> list = new ArrayList<>();
		for (Integer size : map.keySet()) {
			list.add(size);
		}
		long ans = 0;
		for (int i = 0; i < list.size(); i++) {
			int size = list.get(i);
			int num1 = map.get(size);
			ans += 1L * (num1 - 1) * num1 * size * size / 2;
			for (int j = i + 1; j < list.size(); j++) {
				int size2 = list.get(j);
				int num2 = map.get(size2);
				ans += 1L * num1 * num2 * size * size2;
			}
		}
		return ans;
	}

	@Override
	public void debug4() {
		H.compare(12, this.countPairs(6, H.str2array1("[[0,1],[2,3],[4,5]]")));
		H.compare(4999950000l, this.countPairs(100000, H.str2array1("[]")));
		H.compare(14, this.countPairs(7, H.str2array1("[[0,2],[0,5],[2,4],[1,6],[5,4]]")));

	}

	@Override
	public void debug3() {
		// TODO Auto-generated method stub

	}

	@Override
	public void debug2() {
		// TODO Auto-generated method stub

	}

	@Override
	public void debug1() {
		// TODO Auto-generated method stub

	}

	public static void main(String[] args) {
		Solution so = new Solution();
		so.debug1();
		so.debug2();
		so.debug3();
		so.debug4();

	}

}
